import taichi as ti
import numpy as np
# 选择cpu还是gpu
ti.init(arch=ti.cpu)  # Try to run on GPU
# 计算品质，越大算得越准确，但也越慢。
quality = 2  # Use a larger value for higher-res simulations
# 粒子数量，网格数量
n_particles, n_grid = 9000 * quality ** 2, 128 * quality
# 每个网格宽度ΔX=1/128，以及它的倒数inv_dx
dx, inv_dx = 1 / n_grid, float(n_grid)
# deltaTime，演算的时间间隔
dt = 1e-4 / quality
# 体积vol，密度 (rho就是ρ)
p_vol, p_rho = (dx * 0.5)**2, 1
# 质量 = 体积 * 密度
p_mass = p_vol * p_rho
# 以下是材料系数
# E=0.1e4=1000, 泊松分布系数nu=0.2
E, nu = 0.1e4, 0.2  # Young's modulus and Poisson's ratio
# Lame系数，定义材料性质的，分别是μ和λ
mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / \
    ((1+nu) * (1 - 2 * nu))  # Lame parameters
# position位置, 每个粒子有一个位置
x = ti.Vector.field(2, dtype=float, shape=n_particles)  # position
# velocity速度，每个粒子有一个速度
v = ti.Vector.field(2, dtype=float, shape=n_particles)  # velocity
# affine速度场，每个粒子对应一个2x2矩阵
C = ti.Matrix.field(2, 2, dtype=float, shape=n_particles)
# deformation gradient变形梯度矩阵
F = ti.Matrix.field(2, 2, dtype=float, shape=n_particles)
# material id，这个例子里有3种材料，分别是0、1、2
material = ti.field(dtype=int, shape=n_particles)  # material id
# plastic deformation 塑性变形，不可恢复变形
Jp = ti.field(dtype=float, shape=n_particles)  # plastic deformation
# grid node momemtum/velocity 一个128x128矩阵，每个元素是一个Vector2。每个格子一个总体速度
grid_v = ti.Vector.field(2, dtype=float, shape=(
    n_grid, n_grid))  # grid node momentum/velocity
# grid node mass格子质量。128x128矩阵，每个元素是一个数字，每个格子一个质量
grid_m = ti.field(dtype=float, shape=(n_grid, n_grid))  # grid node mass


# 下面的函数是每帧更新用的，先跳过这个函数看主程序初始化~~
# 留到最后看这个函数
@ti.kernel
def substep():
    # 每次都要将每个格子的总速度和总质量归0
    for i, j in grid_m:
        grid_v[i, j] = [0, 0]
        grid_m[i, j] = 0
    # 更新每一个粒子，并让它们属于某一个格子(Particle to Grid)
    for p in x:  # Particle state update and scatter to grid (P2G)
        # 用粒子坐标x[p]，计算出它所属的格子的左下角
        base = (x[p] * inv_dx - 0.5).cast(int)
        # fx是该粒子距离格子左下角的局部坐标
        fx = x[p] * inv_dx - base.cast(float)
        # 二次核函数
        # Quadratic kernels  [http://mpm.graphics   Eqn. 123, with x=fx, fx-1,fx-2]
        # w是一个延迟计算的表达式，无法直接看到它的值，后面很多参数都是类似的
        # w是位置的函数。
        # w是一个权重，会影响当前格子与周围格子的质量、速度
        # 猜测：一个格子边缘的粒子，显然对旁边格子的影响更大
        w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1)
             ** 2, 0.5 * (fx - 0.5) ** 2]
        #每帧更新F，F = somefunc(F, C, dt)
        F[p] = (ti.Matrix.identity(float, 2) + dt * C[p]
                ) @ F[p]  # deformation gradient update
        # h "加硬"系数（雪积累在一起变硬），是根据塑性变形参数Jp算出来的
        # Hardening coefficient: snow gets harder when compressed
        h = ti.exp(10 * (1.0 - Jp[p]))
        # jelly 果冻的加硬系数固定0.3
        if material[p] == 1:  # jelly, make it softer
            h = 0.3
        # μ和λ的值根据h确定，由于加硬，要增大lame系数
        mu, la = mu_0 * h, lambda_0 * h
        # 液体的μ固定为0
        if material[p] == 0:  # liquid
            mu = 0.0
        # ---------------硬核计算开始，按理说不要随意改动-----------------
        U, sig, V = ti.svd(F[p])
        J = 1.0
        for d in ti.static(range(2)):
            new_sig = sig[d, d]
            if material[p] == 2:  # Snow
                new_sig = min(max(sig[d, d], 1 - 2.5e-2),
                              1 + 4.5e-3)  # Plasticity
            Jp[p] *= sig[d, d] / new_sig
            sig[d, d] = new_sig
            J *= new_sig
        if material[p] == 0:  # Reset deformation gradient to avoid numerical instability
            F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J)
        elif material[p] == 2:
            # Reconstruct elastic deformation gradient after plasticity
            F[p] = U @ sig @ V.transpose()
        # (Fixed)Corotated：
        stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose() + \
            ti.Matrix.identity(float, 2) * la * J * (J - 1)
        stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
        affine = stress + p_mass * C[p]
        # ===============硬核计算结束==================
        # 遍历相邻8+1个格子，把粒子参数算到到格子上
        # Loop over 3x3 grid node neighborhood
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            dpos = (offset.cast(float) - fx) * dx
            weight = w[i][0] * w[j][1]
            # 每个粒子的速度、质量叠加，得到当前格子与周围格子的速度、质量
            grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
            grid_m[base + offset] += weight * p_mass
    #遍历所有的格子
    for i, j in grid_m:
        #如果这块格子有质量才需要计算
        if grid_m[i, j] > 0:  # No need for epsilon here
            # 动量转为速度。格子质量越大，格子速度变得越小
            grid_v[i, j] = (1 / grid_m[i, j]) * \
                grid_v[i, j]  # Momentum to velocity
            # 重力影响
            grid_v[i, j][1] -= dt * 50  # gravity
            # 碰到四周墙壁，格子速度强行置为0。数字3就是第几个格子算墙壁的意思，可以改大试试
            if i < 3 and grid_v[i, j][0] < 0:
                grid_v[i, j][0] = 0  # Boundary conditions
            if i > n_grid - 3 and grid_v[i, j][0] > 0:
                grid_v[i, j][0] = 0
            if j < 3 and grid_v[i, j][1] < 0:
                grid_v[i, j][1] = 0
            if j > n_grid - 3 and grid_v[i, j][1] > 0:
                grid_v[i, j][1] = 0
    # 最后，把格子的计算还原到粒子上去。遍历所有粒子
    for p in x:  # grid to particle (G2P)
        base = (x[p] * inv_dx - 0.5).cast(int)
        fx = x[p] * inv_dx - base.cast(float)
        w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1.0)
             ** 2, 0.5 * (fx - 0.5) ** 2]
        new_v = ti.Vector.zero(float, 2)
        new_C = ti.Matrix.zero(float, 2, 2)
        # loop over 3x3 grid node neighborhood
        # 同样，要考虑该粒子周围的几个格子
        for i, j in ti.static(ti.ndrange(3, 3)):
            dpos = ti.Vector([i, j]).cast(float) - fx
            g_v = grid_v[base + ti.Vector([i, j])]
            weight = w[i][0] * w[j][1]
            # 新速度 += 权重 * 格子速度
            new_v += weight * g_v
            new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
        # 更新粒子的速度、C
        v[p], C[p] = new_v, new_C
        # 位置 += 速度 * ΔTime
        x[p] += dt * v[p]  # advection

# 分三组，每组一个正方形
group_size = n_particles // 3


@ti.kernel
def initialize():
    for i in range(n_particles):
        # 位置都是归1化的，取值0~1
        x[i] = [ti.random() * 0.2 + 0.3 + 0.10 * (i // group_size),
                ti.random() * 0.2 + 0.05 + 0.32 * (i // group_size)]
        # 0: fluid流体 1: jelly果冻 2: snow雪
        material[i] = i // group_size  # 0: fluid 1: jelly 2: snow
        # 初始速度0
        v[i] = ti.Matrix([0, 0])
        # 变形梯度矩阵初始值
        F[i] = ti.Matrix([[1, 0], [0, 1]])
        # 塑料变形参数初始值
        Jp[i] = 1


initialize()
# # 初始化ti.GUI，显示用
# gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)
# # 运行直到esc
# while not gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT):
#     # 每帧计算若干次
#     for s in range(int(2e-3 // dt)):
#         substep()
#     # 三种颜色，每组一个颜色
#     colors = np.array([0x068587, 0xED553B, 0xEEEEF0], dtype=np.uint32)
#     # 根据位置画小圆点，根据材质颜色3选1
#     gui.circles(x.to_numpy(), radius=1.5, color=colors[material.to_numpy()])
#     # Change to gui.show(f'{frame:06d}.png') to write images to disk
#     # 显示
#     gui.show()
import time
count = 0
last = time.time()
while True:
    substep()
    count+=1
    if count > 10000:
        break
now = time.time()
print(now-last)